# pylint: disable=line-too-long,useless-suppression
# coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------
"""Customize generated code here.

Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
"""

import asyncio  # pylint: disable=do-not-import-asyncio
from typing import Any, Callable, IO, Coroutine, List, Optional, Union, cast

from azure.core.exceptions import ResourceNotFoundError
from azure.core.polling import AsyncLROPoller, AsyncNoPolling, AsyncPollingMethod

from azure.confidentialledger.aio._operations._operations import (
    _ConfidentialLedgerClientOperationsMixin as GeneratedOperationsMixin,
)
from azure.confidentialledger.aio._operations._operations import ClsType, JSON
from azure.confidentialledger._operations._patch import BaseStatePollingMethod
import azure.confidentialledger.models as _models

__all__: List[str] = [
    "_ConfidentialLedgerClientOperationsMixin"
]  # Add all objects you want publicly available to users at this package level


def patch_sdk():
    """Do not remove from this file.

    `patch_sdk` is a last resort escape hatch that allows you to do customizations
    you can't accomplish using the techniques described in
    https://aka.ms/azsdk/python/dpcodegen/python/customize
    """


class AsyncStatePollingMethod(BaseStatePollingMethod, AsyncPollingMethod):
    """Polling method for methods returning responses containing a 'state' field; the polling
    completes when 'state' becomes a desired value.
    """

    def __init__(
        self,
        operation: Callable[[], Coroutine[Any, Any, JSON]],
        desired_state: str,
        polling_interval_s: float,
        retry_not_found: bool,
    ):
        super().__init__(operation, desired_state, polling_interval_s, retry_not_found)

    async def run(self) -> None:
        try:
            while not self.finished():
                try:
                    response = await self._operation()
                    self._evaluate_response(response)
                except ResourceNotFoundError as not_found_exception:
                    # We'll allow some instances of resource not found to account for replication
                    # delay if session stickiness is lost.

                    self._not_found_count += 1

                    not_retryable = not self._retry_not_found or self._give_up_not_found_error(not_found_exception)

                    if not_retryable or self._not_found_count >= 3:
                        raise
                if not self.finished():
                    await asyncio.sleep(self._polling_interval_s)
        except Exception:
            self._status = "failed"
            raise


class _ConfidentialLedgerClientOperationsMixin(GeneratedOperationsMixin):
    async def begin_get_ledger_entry(
        self, transaction_id: str, *, collection_id: Optional[str] = None, **kwargs: Any
    ) -> AsyncLROPoller[_models.LedgerQueryResult]:
        """Returns a poller to fetch the ledger entry at the specified transaction id.

        A collection id may optionally be specified to indicate the collection from which to fetch
        the value.

        To return older ledger entries, the relevant sections of the ledger must be
        read from disk and validated. To prevent blocking within the enclave, the
        response will indicate whether the entry is ready and part of the response, or
        if the loading is still ongoing.

        :param transaction_id: Identifies a write transaction. Required.
        :type transaction_id: str
        :keyword collection_id: The collection id. Default value is None.
        :paramtype collection_id: str
        :return: An instance of AsyncLROPoller that returns a LedgerQueryResult for the ledger entry.
        :rtype: ~azure.core.polling.AsyncLROPoller[~azure.confidentialledger.models.LedgerQueryResult]
        :raises ~azure.core.exceptions.HttpResponseError:
        """
        polling = kwargs.pop("polling", True)  # type: Union[bool, AsyncPollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        async def operation() -> JSON:
            return await super(_ConfidentialLedgerClientOperationsMixin, self).get_ledger_entry(
                transaction_id, collection_id=collection_id, **kwargs
            )

        initial_response = await operation()

        if polling is True:
            polling_method = cast(AsyncPollingMethod, AsyncStatePollingMethod(operation, "Ready", lro_delay, False))
        elif polling is False:
            polling_method = cast(AsyncPollingMethod, AsyncNoPolling())
        else:
            polling_method = polling
        return AsyncLROPoller(self._client, initial_response, lambda x: x, polling_method)

    async def begin_get_receipt(self, transaction_id: str, **kwargs: Any) -> AsyncLROPoller[_models.TransactionReceipt]:
        """Returns a poller for getting a receipt certifying ledger contents at a particular
        transaction id.

        :param transaction_id: Identifies a write transaction. Required.
        :type transaction_id: str
        :return: An instance of AsyncLROPoller that returns a TransactionReceipt for the receipt.
        :rtype: ~azure.core.polling.AsyncLROPoller[~azure.confidentialledger.models.TransactionReceipt]
        :raises ~azure.core.exceptions.HttpResponseError:
        """
        polling = kwargs.pop("polling", True)  # type: Union[bool, AsyncPollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        async def operation() -> JSON:
            return await super(_ConfidentialLedgerClientOperationsMixin, self).get_receipt(
                transaction_id=transaction_id, **kwargs
            )

        initial_response = await operation()

        if polling is True:
            polling_method = cast(AsyncPollingMethod, AsyncStatePollingMethod(operation, "Ready", lro_delay, False))
        elif polling is False:
            polling_method = cast(AsyncPollingMethod, AsyncNoPolling())
        else:
            polling_method = polling
        return AsyncLROPoller(self._client, initial_response, lambda x: x, polling_method)

    async def begin_create_ledger_entry(
        self, entry: Union[_models.LedgerEntry, JSON, IO[bytes]], *, collection_id: Optional[str] = None, **kwargs: Any
    ) -> AsyncLROPoller[_models.TransactionStatus]:
        """Writes a ledger entry and returns a poller to wait for it to be durably committed. The
        poller returns the result for the initial call to create the ledger entry.

        A collection id may optionally be specified.

        :param entry: Ledger entry. Required.
        :type entry: ~azure.confidentialledger.models.LedgerEntry or JSON or IO[bytes]
        :keyword collection_id: The collection id. Default value is None.
        :paramtype collection_id: str
        :return: AsyncLROPoller[TransactionStatus]. The TransactionStatus is compatible with MutableMapping
        :rtype: ~azure.core.polling.AsyncLROPoller[~azure.confidentialledger.models.TransactionStatus]
        :raises ~azure.core.exceptions.HttpResponseError:
        """

        # Pop arguments that are unexpected in the pipeline.

        polling = kwargs.pop("polling", True)  # type: Union[bool, AsyncPollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        # Pop the custom deserializer, if any, so we know the format of the response and can
        # retrieve the transactionId. Serialize the response later.

        cls = kwargs.pop("cls", None)  # type: ClsType[JSON]
        kwargs["cls"] = lambda pipeline_response, json_response, headers: (
            pipeline_response,
            {
                **json_response,
                "transactionId": headers.get("x-ms-ccf-transaction-id") if headers else None,
            },
            headers,
        )

        post_pipeline_response, post_result, post_headers = await self.create_ledger_entry(
            entry, collection_id=collection_id, **kwargs
        )

        # Delete the cls because it should only apply to the create_ledger_entry response, not the
        # wait_for_commit call.

        del kwargs["cls"]

        transaction_id = post_result["transactionId"]  # type: ignore

        kwargs["polling"] = polling
        kwargs["polling_interval"] = lro_delay

        if cls:
            kwargs["_create_ledger_entry_response"] = cls(
                post_pipeline_response, cast(JSON, post_result), post_headers  # type: ignore
            )
        else:
            kwargs["_create_ledger_entry_response"] = post_result
        return await self.begin_wait_for_commit(transaction_id, **kwargs)

    async def begin_wait_for_commit(
        self, transaction_id: str, **kwargs: Any
    ) -> AsyncLROPoller[_models.TransactionStatus]:
        """Creates a poller that queries the state of the specified transaction until it is
        Committed, a state that indicates the transaction is durably stored in the Confidential
        Ledger.

        :param transaction_id: Identifies a write transaction. Required.
        :type transaction_id: str
        :return: An instance of AsyncLROPoller returning a TransactionStatus object describing the transaction status.
        :rtype: ~azure.core.polling.AsyncLROPoller[~azure.confidentialledger.models.TransactionStatus]
        :raises ~azure.core.exceptions.HttpResponseError:
        """
        polling = kwargs.pop("polling", True)  # type: Union[bool, AsyncPollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        # If this poller was called from begin_create_ledger_entry, we should return the
        # create_ledger_entry response, not the transaction status.

        post_result = kwargs.pop("_create_ledger_entry_response", None)

        def deserialization_callback(x):
            return x if post_result is None else post_result

        async def operation() -> JSON:
            return await super(_ConfidentialLedgerClientOperationsMixin, self).get_transaction_status(
                transaction_id=transaction_id, **kwargs
            )

        try:
            initial_response = await operation()
        except ResourceNotFoundError:
            if polling is False or polling is None:
                raise
            # This method allows for temporary resource not found errors, which may occur if session
            # stickiness is lost and there is replication lag.

            initial_response = {}
        if polling is True:
            polling_method = cast(
                AsyncPollingMethod,
                AsyncStatePollingMethod(operation, "Committed", lro_delay, True),
            )
        elif polling is False:
            polling_method = cast(AsyncPollingMethod, AsyncNoPolling())
        else:
            polling_method = polling
        return AsyncLROPoller(self._client, initial_response, deserialization_callback, polling_method)

    async def create_ledger_entry(
        self, entry: Union[_models.LedgerEntry, JSON, IO[bytes]], *, collection_id: Optional[str] = None, **kwargs: Any
    ) -> _models.LedgerWriteResult:
        """Writes a ledger entry.

        A collection id may optionally be specified.

        :param entry: Ledger entry. Is one of the following types: LedgerEntry, JSON, IO[bytes]
         Required.
        :type entry: ~azure.confidentialledger.models.LedgerEntry or JSON or IO[bytes]
        :keyword collection_id: The collection id. Default value is None.
        :paramtype collection_id: str
        :return: LedgerWriteResult. The LedgerWriteResult is compatible with MutableMapping
        :rtype: ~azure.confidentialledger.models.LedgerWriteResult
        :raises ~azure.core.exceptions.HttpResponseError:

        Example:
            .. code-block:: python

                # JSON input template you can fill out and use as your body input.
                entry = {
                    "collectionId": {
                        "collectionId": "str"  # Required.
                    },
                    "contents": "str",  # Required. Contents of the ledger entry.
                    "transactionId": "str"  # Optional. A unique identifier for the state of the
                      ledger. If returned as part of a LedgerEntry, it indicates the state from which
                      the entry was read.
                }
        """

        kwargs["cls"] = kwargs.get(
            "cls",
            lambda _, json_response, headers: {
                **json_response,
                "transactionId": headers.get("x-ms-ccf-transaction-id") if headers else None,
            },
        )
        return await super().create_ledger_entry(entry, collection_id=collection_id, **kwargs)
